// "Copyright [2021] <Copyright Shengkai Lin>"

// The topology of this network is shown as the ./oxc_cluster.png
// TODO(Lin) : The flexibility is still a problem

package cluster.networks;

import cluster.common.*;
import cluster.dispatcher.*;
import cluster.network_controller.*;
import cluster.network_controller.topo_manager.*;
import cluster.network_devices.*;
import cluster.optical_devices.*;
import cluster.resource_scheduler.*;
import cluster.task_scheduler.*;
import inet.networklayer.configurator.ipv4.*;
import inet.networklayer.ipv4.*;

network OXCNetwork
{
    parameters:
        int numMachine = default(2);
        int numOxc = default(64);
        int numOpticalswitch8x64 = default(numMachine);  // Default it is the same as the AIServer num
        int numGpusperMachine = default(8);
        int numNvlink_port = default(6);
        double switch_delay = default(0);
        // enum ClusterType {
        //    OpticalCluster,
        //    ElectricalCluster,
        //    IbCluster
        //};        
        int clustertype = default(0);
        int inputType = default(0);
        int maxnumBatches = default(1);
        int modelsize_ = default(1000000);

    submodules:
        oxc[numOxc]: OpticalDevice{
            parameters:
                opticaldevice_type=0;
                numOptical_port=numMachine;
                my_index = index;
                switch_delay = switch_delay;
        }

        opticalswitch[numMachine] : OpticalDevice{
            parameters:
                opticaldevice_type=1;
                numOptical_port= numGpusperMachine + numOxc;
                my_index = index;
                switch_delay = 0.00000005;
        }

        AIServer[numMachine] : Machine{
            parameters:
                numGpus = numGpusperMachine;
                numNvlink_port = numNvlink_port;
                my_index = index;
        }

        topo_manager : TopoManager {}

        dispatcher : Dispatcher {}

        resource_scheduler : ResourceScheduler {}

        task_scheduler : TaskScheduler {
        parameters:
            maxnumGpus = numGpusperMachine * numMachine;
            modelsize = modelsize_;
            maxnumBatches = maxnumBatches;
        }

        network_controller : NetworkController {}

        configurator: Ipv4NetworkConfigurator {
            parameters:
                @display("p=100,100;is=s");
                dumpAddresses = true;
                dumpTopology = true;
                dumpLinks = true;
                //dumpConfig = "./networkconfig.xml";
                // addSubnetRoutes = false;
                // addDefaultRoutes = false;
                // optimizeRoutes = false;
                // config = xml("<config><interface hosts='**.nic[*]' address='192.168.x.x' netmask='255.225.0.0'/></config>");
                // config = xmldoc("/ip.xml");
                assignAddresses = false;
                addSubnetRoutes = false;
                optimizeRoutes = false;
                networkType = 0;
        }

    connections allowunconnected:
        // Connect the data plane
        // Connect the machines first
        for i=0..numMachine-1, for j=0..numGpusperMachine-1 {
            AIServer[i].out_port[j] <-->Fiber<--> opticalswitch[i].optical_port[j];
        }
        for i=0..numMachine-1, for j=0..numOxc-1 {
            opticalswitch[i].optical_port[j+numGpusperMachine]<-->FiberHalf<-->oxc[j].optical_port[i];
        }

        // Connect the task generator and controller
        task_scheduler.port --> dispatcher.task_generator_port;

        // Connect the controller - AIServer
        for i=0..numMachine-1 {
            dispatcher.machine_port++<-->ControlPlaneLink<-->AIServer[i].controller_port;
        }

        dispatcher.network_routing_port<--ControlPlaneLink<--network_controller.dispatcher_port$o;

        // Connect the switch - controller
        for i=0..numOxc-1 {
            network_controller.switch_port++ <-->ControlPlaneLink<--> oxc[i].controller_port;
        }
        for i=0..numOpticalswitch8x64-1 {
            network_controller.switch_port++ <-->ControlPlaneLink<--> opticalswitch[i].controller_port;
        }
}